from mindspore import nn, ops, Tensor
import numpy as np
import mindspore as ms

class LightEyes(nn.Cell):
    
    def __init__(self):
        super().__init__()
        self.conv1 = nn.SequentialCell(
            # conv1_1
            nn.Conv2d(in_channels=3, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        self.conv2_12 = nn.SequentialCell(
            # conv1_2
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_3
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_4
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_5
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_6
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_7
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_8
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_9
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_10
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_11
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            # conv1_12
            nn.Conv2d(in_channels=16, out_channels=16,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        self.conv13 = nn.SequentialCell(
            # conv1_13
            nn.Conv2d(in_channels=16, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        self.conv14_15 = nn.SequentialCell(
            #conv1_14
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_15
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        self.conv16_24 = nn.SequentialCell(
            #conv1_16
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_17
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_18
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_19
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_20
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_21
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_22
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_23
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_24
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv25_27 = nn.SequentialCell(
            #conv1_25
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_26
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU(),
            #conv1_27
            nn.Conv2d(in_channels=8, out_channels=8,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform'),
            nn.ReLU()
        )
        #conv1_28
        self.conv28 = nn.Conv2d(in_channels=8, out_channels=1,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform')
        #conv1_29
        #score1
        self.conv29 = nn.Conv2d(in_channels=8, out_channels=1,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform')
        #conv1_30
        #score2
        self.conv30 = nn.Conv2d(in_channels=8, out_channels=1,
                    pad_mode="pad", padding=1,
                    has_bias=True, 
                    kernel_size=3, weight_init='xavier_uniform')
        #conv1_31
        #loss
        self.conv31 =nn.Conv2d(in_channels=3, out_channels=1,
                        pad_mode="pad", padding=1,
                        has_bias=True, 
                        kernel_size=3, weight_init='xavier_uniform')
        #sigmoid after conv31
        self.sigmoid1=nn.Sigmoid()
        #deconv1
        #up2
        #score3
        self.deconv1 = nn.Conv2dTranspose(in_channels=1, out_channels=1,
                    has_bias=True, 
                    kernel_size=4, weight_init='xavier_uniform',
                    stride=2)
    
    def construct(self, x):
        x1 = self.conv1(x)
        x2 = self.conv2_12(x1)
        x3 = self.conv13(x2)
        x4 = self.conv14_15(x3)
        score1 = self.conv29(x4)
        x5 = self.conv16_24(x4)
        score2 = self.conv30(x5)
        p1 = self.pool1(x5)
        x6 = self.conv25_27(p1)
        x7 = self.conv28(x6)
        x8 = self.deconv1(x7)#尺寸尚不是原始图像尺寸
        cropsize = x.shape[-2:]
        x8_nhwc = ops.transpose(x8, (0, 2, 3, 1))
        score3_nhwc = ops.crop_and_resize(x8_nhwc, 
                        Tensor([[0,0,1,1]], dtype=ms.float32), 
                        Tensor([0], dtype=ms.int32), 
                        cropsize)
        score3 = ops.transpose(score3_nhwc, (0, 3, 1, 2))
        x9_cat = ops.cat((score1, score2, score3), 1)
        out = self.conv31(x9_cat)
        probability=self.sigmoid1(out)#分割概率矩阵
        return probability,out
        